import os
import torch
import dataLoad
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torch.utils.data.dataset as Dataset

import MIEEGModel

# 超参数
batch_size = 64
epochs = 100
# cuda
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 设置seed
seed = 42
np.random.seed(seed)
torch.manual_seed(seed)

if __name__ == '__main__':
    # 定义数据集存放目录
    """
    # 查看数据信息
    file_path = './BCICIV_2b_gdf/B0101T.gdf'
    check_data_info(file_path, img_save=True, img_dir='./img')
    """
    file_dir = './BCICIV_2b_gdf/train_data'
    data, labels = dataLoad.load_data_BCICIV_2b_gdf(file_dir)
    # 构造数据集
    dataset = dataLoad.EEGDataset(data, labels)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    # dataLoad.print_version()
    MIEEGModel.model_train(dataloader, device, epochs, model_save=True)

